#include "util.h"

#include <iostream>
#include <fstream>
#include <sstream>
#include <strings.h>
#include <math.h>
#include <assert.h>
#include <map>
#include <vector>
#include <set>

// map data for classes
static std::map<std::string,int> class_map;
static std::map<int,std::string> rev_class_map;
static std::map<std::string,std::set<int> > parent_map;

typedef std::map<std::string,int>::value_type mapVal;
typedef std::map<int,std::string>::value_type revMapVal;
typedef std::map<std::string,std::set<int> >::value_type parMapVal;

static unsigned int class_counter = 0;

bool read_header(std::istream &, Parameter&);
bool read_class(std::istream&, int&);
bool load_data_svmlight(char *filename, SparseArray **&X, unsigned int *&Y, Parameter &param);
bool _load_string_data(char *filename, std::string *&X, unsigned int *&Y, int &character_count,
                       Parameter &param);

std::map<std::string,std::set<int> > &parent_class_sets()
{
    return parent_map;
}

bool load_multi_data(char *vector_file, char *string_file, char *trees_file,  
                     MultiDataSet *&data, Parameter &param)
{
//    std::cout << "--- Loading data ---" << std::endl;
    // make sure we have some data
    assert(vector_file != 0 || string_file != 0 || trees_file != 0);

    std::string *sX=0;
    SparseArray **vX=0;
    SparseArray **tX=0;
    unsigned int *sY=0, *vY=0, *tY=0; 

    int character_count=0;
    if (string_file != 0)
        if(!_load_string_data(string_file, sX, sY, character_count, param)) return false;
//   class_counter = 0;
    if (vector_file != 0)
        if(!load_data_svmlight(vector_file, vX, vY, param)) return false;
//    class_counter = 0;
//  if (trees_file != 0)
//      if(!load_tree_data(trees_file, tX, tY, param)) return false;

    // verify the labels match
//  if (vector_file != 0 && string_file != 0)
//  {
//      for (unsigned int i=0; i < param.m; ++i)
//      {
//          if(vY[i] != sY[i]);
//              std::cout << i << " " << vY[i] << " " << sY[i] << std::endl;
//          assert(vY[i] == sY[i]);
//      }
//      delete sY;
//  }

    if (vector_file != 0)
    {
        data = new MultiDataSet(vX, sX, tX, vY, param.m);
        delete sY;
        delete tY;
    }
    else if (string_file != 0)
    {
        data = new MultiDataSet(vX, sX, tX, sY, param.m);
        delete tY;
    }
    else
        data = new MultiDataSet(vX, sX, tX, tY, param.m);

    data->set_chars(character_count);
//
//  std::map<std::string,std::set<int> >::iterator p_it = parent_map.begin();
//  for(; p_it != parent_map.end(); ++p_it)
//  {
//      std::cout << p_it->first;
//      std::set<int>::iterator v_it = p_it->second.begin();
//      for (; v_it != p_it->second.end(); ++v_it)
//          std::cout << " " << *v_it;
//      std::cout << std::endl;
//  }
//
//    std::cout << "--- Finished loading data ---" << std::endl;
    return true;
}

bool _load_string_data(char *filename, std::string *&X, unsigned int *&Y, int &character_count,
                       Parameter &param)
{
    // open the data file for reading
    std::ifstream inFile(filename);
    if(!inFile)
    {
        std::cerr << "unable to open input file: "
            << filename << "\n";
        return false;
    }

    std::cout << "Loading string instance format events file...";

    // do a first pass to get the data dimensions
    int i=0; // number of events
    character_count=0;
    int cls=0;
    std::string buffer;

    while(!inFile.eof())
    {
        if(!read_class(inFile, cls))
            break;
        std::getline(inFile, buffer, '\n');
        character_count += (buffer.length()+1);
        i++;
    } 
    inFile.close();

    param.m = i;
    param.k = class_counter;

    // create the actual events data structures and read in the events
    Y = new unsigned int [param.m];
    X = new std::string[param.m];

    std::ifstream fd(filename);
    for(int i=0; !fd.eof(); ++i)
    {
        if(!read_class(fd, cls))
            break;
        Y[i] = cls; 

        std::getline(fd, buffer, '\n');
        for(unsigned int j=0; j < buffer.length(); ++j)
        {
            if ((int)buffer[j] < 0)
            {
                std::cout << buffer << std::endl;
                assert(false);
            }
        }
        X[i] = buffer;
    } 
    fd.close();
    std::cout << "done" << std::endl;
    return true;
}

bool load_data_svmlight(char *filename, SparseArray **&X, unsigned int *&Y, Parameter &param)
{
    // open the data file for reading
    std::ifstream inFile(filename);
    if(!inFile)
    {
        std::cerr << "unable to open input file: "
            << filename << "\n";
        return false;
    }

    std::cout << "Loading SVMLight format events file...";

    // do a first pass to get the data dimensions
    char tmp_char = 0;
    int i=0; // number of events
    param.n = 0; // number of features
    unsigned int index;
    unsigned int num_missing_labels=0;
    int cls;
    double value;
    std::string buffer;

    while(!inFile.eof())
    {
        if(!read_class(inFile, cls))
            break;
        if (cls < 0) num_missing_labels++;

        std::getline(inFile, buffer, '\n');
        std::istringstream ss(buffer);
        int prev_index = -1;
        while(ss)
        {
            ss >> index;
            if(!ss) break;

            ss.get(tmp_char);
            assert(tmp_char == ':');
            ss >> value;
            if(index > param.n)
                param.n = index;

            // index values must be sorted
//            if((int)index <= prev_index)
//                std::cout << (int)index << '-' << prev_index << std::endl;
//            assert((int)index > prev_index);
            prev_index = index;
        }
        i++;
    } 
    inFile.close();

    param.n += 1;
    param.m = i;
    param.k = class_counter;
    param.u = num_missing_labels;
    param.basis = param.m-param.u;

    // create the actual events data structures and read in the events
    Y = new unsigned int [param.m];
    X = new SparseArray*[param.m];
    SparseArray *sa=0;

    std::ifstream fd(filename);
    for(int i=0; !fd.eof(); ++i)
    {
        if(!read_class(fd, cls))
            break;
        Y[i] = cls; 

        std::getline(fd, buffer, '\n');
        std::istringstream ss(buffer);
        sa = new SparseArray(param.n);
        std::vector< std::pair<int,double> > vec;
        while(ss)
        {
            ss >> index;
            if(!ss) break;

            ss.get(tmp_char);
            assert(tmp_char == ':');
            ss >> value;
            vec.push_back( std::pair<int,double>(index, value));
        }
        std::sort(vec.begin(), vec.end());
        for(std::vector<std::pair<int,double> >::iterator it=vec.begin(); it != vec.end(); ++it)
            sa->insert(it->first, it->second);
        X[i] = sa;
    } 
    fd.close();
    std::cout << param.m << " instances loaded" << std::endl;

    return true;
}

bool read_class(std::istream &in, int &cls)
{
    std::string tag;

    do { in >> tag; }
    while(!in.eof() && tag.size() == 0);
    
    if(in.eof())
        return false;

    if(class_map.count(tag) == 0)
    {
        int cls_val=-1;
        if (tag != "NONE")
            cls_val = class_counter++;

        rev_class_map.insert(revMapVal(cls_val,tag)); 
        class_map.insert(mapVal(tag,cls_val)); 
    }

    cls = class_map[tag];

    unsigned int parent_index=0;
    if((parent_index == tag.find(':')) != std::string::npos)
    {
        std::string parent = tag.substr(0,parent_index);
        parent_map[parent].insert(cls);
    }

    return true;
}

std::string parent_class(const std::string &tag)
{
    std::string result;
    unsigned int parent_index=0;
    if((parent_index == tag.find(':')) != std::string::npos)
        return tag.substr(0,parent_index);
    return tag;        
}

std::string map_class(unsigned int cls)
{
    return rev_class_map[cls];
}

// select heldout datasets
bool get_heldout_set(int fold, MultiDataSet *data, MultiDataSet *&train_data, 
                     MultiDataSet *&heldout_data, Parameter *param)
{
    int heldout_size = param->m / param->heldout_sets;
    int train_counter = 0, heldout_counter = 0, train_character_count=0, heldout_character_count=0;

    // allocate the data structures
    unsigned int *train_classes = new unsigned int [param->m - heldout_size];
    std::string *train_strings=0;
    SparseArray **train_vectors=0;
    SparseArray **train_trees=0;
    if (data->got_strings)
        train_strings = new std::string[param->m - heldout_size];
    if (data->got_vectors)
        train_vectors = new SparseArray*[param->m - heldout_size];
    if (data->got_trees)
        train_trees = new SparseArray*[param->m - heldout_size];

    unsigned int *heldout_classes = new unsigned int [heldout_size];
    std::string *heldout_strings=0;
    SparseArray **heldout_vectors=0;
    SparseArray **heldout_trees=0;
    if (data->got_strings)
        heldout_strings = new std::string[heldout_size];
    if (data->got_vectors)
        heldout_vectors = new SparseArray*[heldout_size];
    if (data->got_trees)
        heldout_trees = new SparseArray*[heldout_size];

    // copy the events and classes into the appropriate places
    for(unsigned int i=0; i < param->m; ++i)
    {
        if((i/heldout_size) != (unsigned int)fold)
//        if((i % heldout_size) != (unsigned int)fold)
        {
            if (data->got_strings)
                train_strings[train_counter] = data->string(i);    
            if (data->got_vectors)
                train_vectors[train_counter] = data->vector(i);    
            if (data->got_trees)
                train_trees[train_counter] = data->tree(i);    
            train_classes[train_counter++] = data->label(i);    
            if (data->got_strings)
                train_character_count += (data->string(i).length() + 1);
        }
        else
        {
            if (data->got_strings)
                heldout_strings[heldout_counter] = data->string(i);    
            if (data->got_vectors)
                heldout_vectors[heldout_counter] = data->vector(i);    
            if (data->got_trees)
                heldout_trees[heldout_counter] = data->tree(i);    
            heldout_classes[heldout_counter++] = data->label(i);    
            if (data->got_strings)
                heldout_character_count += (data->string(i).length() + 1);
        }
    }
    train_data = new MultiDataSet(train_vectors, train_strings, train_trees, train_classes, 
                                  param->m-heldout_size);
    heldout_data = new MultiDataSet(heldout_vectors, heldout_strings, heldout_trees,
                                  heldout_classes, heldout_size);

    if (data->got_strings)
    {
        train_data->set_chars(train_character_count);
        heldout_data->set_chars(heldout_character_count);
    } 
    return true;
}

///////////// Sequence Data ////////////////
MultiDataSet *SequenceDataSet::flatten()
{
    // count instances
    int instance_counter = 0;
    for (int i=0; i < m_num_sequences; ++i)
        instance_counter += m_sequences[i].size();

    // create multi data
    unsigned int *sY = new unsigned int [instance_counter];
    SparseArray **sX = new SparseArray* [instance_counter];
    instance_counter = 0;
    for (int i=0; i < m_num_sequences; ++i)
        for (unsigned int j=0; j < m_sequences[i].size(); ++j)
        {
//            sX[instance_counter] = m_sequences[i][j];
            sX[instance_counter] = add_context(i,j,1);
            sY[instance_counter] = m_labels[i][j];
            instance_counter++;
        }
    return new MultiDataSet(sX, 0, 0, sY, instance_counter);
}

SparseArray* SequenceDataSet::add_context(int sequence, int instance, int order)
{
    return add_context(sequence, instance, m_labels[sequence], order);
}

SparseArray* SequenceDataSet::add_context(int sequence, int instance, std::vector<unsigned int> &labels, 
                                          int order)
{
    assert(order > 0);
    assert(sequence < m_num_sequences);
    assert((int)m_sequences[sequence].size() > instance);
    SparseArray *target = new SparseArray(*m_sequences[sequence][instance]);

    int original_length = target->length();
    target->set_length(original_length + order*(m_klasses+1));
    for (int o=1; o <= order; ++o)
    {
        if((instance-o) < 0)
            target->insert(original_length+(o-1)*(m_klasses+1)+m_klasses, 1.0);
        else
            target->insert(original_length+(o-1)*(m_klasses+1)+labels[instance-o], 1.0);
    }
    return target;
}

SparseArray* SequenceDataSet::add_context(SparseArray *instance, std::vector<unsigned int> &labels, 
                         unsigned int klasses, int order)
{
    SparseArray *target = new SparseArray(*instance);

    int original_length = target->length();
    int labels_length = labels.size();
    target->set_length(original_length + order*(klasses+1));
    for (int o=1; o <= order; ++o)
    {
        if((labels_length-o) < 0)
            target->insert(original_length+(o-1)*(klasses+1)+klasses, 1.0);
        else
            target->insert(original_length+(o-1)*(klasses+1)+labels[labels_length-o], 1.0);
    }
    return target;
}
SparseArray* SequenceDataSet::add_context(SparseArray *instance, int label, unsigned int klasses)
{
    SparseArray *target = new SparseArray(*instance);

    int original_length = target->length();
    target->set_length(original_length + klasses+1);

    if(label < 0)
        target->insert(original_length+klasses, 1.0);
    else
        target->insert(original_length+label, 1.0);
    return target;
}

bool load_sequence_data(char *sequence_file, SequenceDataSet *&data, Parameter &param, bool no_add)
{
    std::cout << "--- Loading Sequences ---" << std::endl;

    std::vector<SparseArray*> *sX=0;
    std::vector<unsigned int> *sY=0; 

    // open the data file for reading
    std::ifstream inFile(sequence_file);
    if(!inFile)
    {
        std::cerr << "unable to open input file: " << sequence_file<< "\n";
        return false;
    }

    std::cout << "Loading sequence events file...";

    // do a first pass to get the data dimensions
    char tmp_char = 0;
    int i=0; // number of events
    int si=0; // number of sequences
    if (!no_add)
        param.n = 0; // number of features
    unsigned int index;
    int cls;
    double value;
    std::string buffer;

    while(!inFile.eof())
    {
        if(inFile.peek() == '\n')
        {
            inFile.ignore();
            si++;
            if (si % 1000 == 0)
                std::cout << "Sequence " << si << std::endl;
            continue;
        }
        if(!read_class(inFile, cls))
            break;

        std::getline(inFile, buffer, '\n');
        std::istringstream ss(buffer);
        int prev_index = -1;
        while(ss)
        {
            ss >> index;
            if(!ss) break;

            ss.get(tmp_char);
            assert(tmp_char == ':');
            ss >> value;
            if(!no_add && index > param.n)
                param.n = index;
            prev_index = index;
        }
        i++;
    } 
    inFile.close();

    if (!no_add)
        param.n += 1;
    param.s = si;
    param.m = i;
    param.k = class_counter;

    // create the actual events data structures and read in the events
    sY = new std::vector<unsigned int> [param.s];
    sX = new std::vector<SparseArray*> [param.s];
    SparseArray *sa=0;

    std::ifstream fd(sequence_file);
    si = 0;
    for(int i=0; !fd.eof(); ++i)
    {
        if(fd.peek() == '\n')
        {
            fd.ignore();
            si++;
            continue;
        }
        if(!read_class(fd, cls))
            break;
        sY[si].push_back(cls); 

        std::getline(fd, buffer, '\n');
        std::istringstream ss(buffer);
        sa = new SparseArray(param.n);
        std::vector< std::pair<int,double> > vec;
        while(ss)
        {
            ss >> index;
            if(!ss) break;

            ss.get(tmp_char);
            assert(tmp_char == ':');
            ss >> value;
            if (!no_add || index < param.n)
                vec.push_back( std::pair<int,double>(index, value));
        }
        std::sort(vec.begin(), vec.end());
        for(std::vector<std::pair<int,double> >::iterator it=vec.begin(); it != vec.end(); ++it)
            sa->insert(it->first, it->second);
        sX[si].push_back(sa);
    } 
    fd.close();
    std::cout << "done" << std::endl;

    data = new SequenceDataSet(sX, sY, param.s, param.k);

    std::cout << "--- Finished loading Sequences---" << std::endl;
    return true;
}

// select heldout datasets
bool get_heldout_sequences(int fold, SequenceDataSet *data, SequenceDataSet *&train_data, 
                           SequenceDataSet *&heldout_data, Parameter *param, GPParameter *ho_param)
{
    int heldout_size = param->s / param->heldout_sets;
    int train_counter = 0, heldout_counter = 0;

    // allocate the data structures
    std::vector<unsigned int> *train_classes 
        = new std::vector<unsigned int> [param->s-heldout_size];
    std::vector<SparseArray*> *train_sequences 
        = new std::vector<SparseArray*> [param->s-heldout_size];

    std::vector<unsigned int> *heldout_classes
        = new std::vector<unsigned int> [heldout_size];
    std::vector<SparseArray*> *heldout_sequences 
        = new std::vector<SparseArray*> [heldout_size];

    // copy the events and classes into the appropriate places
    int num_train_instances = 0;
    int num_test_instances = 0;
    for(unsigned int i=0; i < param->s; ++i)
    {
        if((i/heldout_size) != (unsigned int)fold)
        {
            train_sequences[train_counter] = data->sequence(i);    
            train_classes[train_counter++] = data->label(i);    
            num_train_instances += data->sequence(i).size();
        }
        else
        {
            heldout_sequences[heldout_counter] = data->sequence(i);    
            heldout_classes[heldout_counter++] = data->label(i);    
            num_test_instances += data->sequence(i).size();
        }
    }
    train_data = new SequenceDataSet(train_sequences, train_classes, param->s-heldout_size, param->k);
    heldout_data = new SequenceDataSet(heldout_sequences, heldout_classes, heldout_size, param->k);

    ho_param->s = param->s - heldout_size;
    ho_param->st = heldout_size;
    ho_param->m = num_train_instances;
    ho_param->basis = num_train_instances;
    ho_param->t = num_test_instances;

    return true;
}

